import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

test_data = torchvision.datasets.CIFAR10(
    # 数据集下载路径
    "./dataset",
    download=False,
    # 测试数据集还是训练数据集
    train=False,
    # 图片是PIL格式，需要转换为tensor
    transform=torchvision.transforms.ToTensor()
)

dataloader = DataLoader(test_data, batch_size=64)


class MyMod(nn.Module):
    def __init__(self):
        super(MyMod, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=3, padding=0, stride=1)

    def forward(self, x):
        return self.conv1(x)
# 使用命令（tensorboard.exe --logdir=logs）查看图像
writer = SummaryWriter("../logs")

mymod = MyMod()
print(mymod)
exit()
step = 0
for data in dataloader:
    imgs, targets = data
    output = mymod(imgs)
    # print(imgs.shape)
    writer.add_images("input", imgs, step)
    # print(output.shape)
    # 图片只有3通道，output要展示需要reshape，第一个参数-1代表torch自动计算有batch_size
    output = torch.reshape(output, [-1, 3, 30, 30])
    # print(output.shape)
    writer.add_images("output", output, step)
    step += 1